import requests
import time
from concurrent.futures import ThreadPoolExecutor

# 服务器配置
SERVER_URL = "http://localhost:5000"
TEST_MESSAGE = "请解释一下人工智能的概念"
MODEL_NAME = "DeepSeek-R1"
USE_STREAM = False
SINGLE_TURN = True
NUM_REQUESTS = 10  # 并发请求数量
MAX_WORKERS = 5  # 线程池最大工作线程数

def send_request():
    payload = {
        "message": TEST_MESSAGE,
        "model": MODEL_NAME,
        "stream": USE_STREAM,
        "singleTurn": SINGLE_TURN
    }
    try:
        start_time = time.time()
        response = requests.post(
            f"{SERVER_URL}/chat",
            json=payload,
            stream=USE_STREAM,
            timeout=30
        )
        end_time = time.time()
        elapsed_time = end_time - start_time
        if response.status_code == 200:
            return True, elapsed_time
        else:
            return False, elapsed_time
    except Exception as e:
        return False, 0

def main():
    success_count = 0
    total_time = 0
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = [executor.submit(send_request) for _ in range(NUM_REQUESTS)]
        for future in futures:
            success, elapsed_time = future.result()
            if success:
                success_count += 1
                total_time += elapsed_time

    success_rate = success_count / NUM_REQUESTS * 100
    average_time = total_time / success_count if success_count > 0 else 0

    print(f"并发请求数量: {NUM_REQUESTS}")
    print(f"成功率: {success_rate:.2f}%")
    print(f"平均响应时间: {average_time:.2f} 秒")

if __name__ == "__main__":
    main()